#include "core/main.h"

void GameModelEvaluator::evaluateModel(const Game &game, const ReplayCollection &replays, const GameModel &model,
                                       const GameModelEvaluatorParams &params)
{
    const SpriteCollection &sprites = game.Collection();
    UINT numSprites = sprites.NonFontSpriteCount();

    vector<string> spriteNames;
    for (UINT spriteId = 0; spriteId < numSprites; spriteId++) {
        spriteNames.push_back(sprites.GetSprite(spriteId).name);
    }

    // Add one for the empty "sprite"
    numSprites++;
    spriteNames.push_back("None");

    // confusionMatrices[f] is the confusion matrxix for f-frame lookahead.
    // each confusion matrix is structured confusionMatrix(actual, predicted)
    vector<Grid<int>> confusionMatrices;
    for (int f = 0; f < params.maxLookAhead; f++) {
        confusionMatrices.push_back(Grid<int>(numSprites, numSprites, 0));
    }

    // totalPredictsions[f] and totalCorrect[f] refer to an f-frame lookahead
    vector<int> totalPredictions(params.maxLookAhead, 0);
    vector<int> totalCorrect(params.maxLookAhead, 0);

    // Grid<int> confusionMatrix(numSprites, numSprites, 0);
    // confusionMatrix(actual, predicted)

    // int totalPredictions = 0;
    // int totalCorrect = 0;

    int frameCount = 0;
    for (auto rit = replays.Replays().begin(); rit != replays.Replays().end(); ++rit) {
        const Replay& replay = **rit;
        for (auto fit1 = replay.Frames().begin(); fit1 != replay.Frames().end(); fit1++) {
            if (frameCount % 100 == 0) cout << "Starting frame " << frameCount << endl;
            frameCount++;

            auto fit2 = fit1;
            fit2++;

            Frame curGuess = *fit1;
            for (int lookAhead = 1; lookAhead <= params.maxLookAhead && fit2 != replay.Frames().end(); lookAhead++) {
                const Frame &realNextFrame = *fit2;
                fit2++;

                // Predict the next frame using the predicion of the current frame
                vector<const Frame *> history;
                history.push_back(&curGuess);
                Frame nextGuess;
                model.AdvanceFrame(history, nextGuess);

                // Keep the train rolling
                curGuess = nextGuess;
                curGuess.controller = realNextFrame.controller;

                FrameGrid realNextGrid(realNextFrame, params.gridScale);
                FrameGrid predNextGrid(nextGuess, params.gridScale);

                for (int y = 0; y < realNextGrid.height() && y < predNextGrid.height(); y++) {
                    for (int x = 0; x < realNextGrid.width() && x < predNextGrid.width(); x++) {
                        int realId = realNextGrid.getId(x, y);
                        int predId = predNextGrid.getId(x, y);

                        // Rejigger IDs in the case that either is empty
                        if (realId == -1) realId = numSprites - 1;
                        if (predId == -1) predId = numSprites - 1;

                        confusionMatrices[lookAhead - 1](realId, predId) = confusionMatrices[lookAhead - 1](realId, predId) + 1;
                        totalPredictions[lookAhead - 1]++;
                        if (realId == predId) totalCorrect[lookAhead - 1]++;
                    }
                }
            }
        }

        /*
        auto fit2 = replay.Frames().begin();
        fit2 += params.lookAhead;
        while (fit2 != replay.Frames().end()) {
            if (frameCount % 100 == 0) cout << frameCount << endl;
            frameCount++;

            const Frame &curFrame = *fit1;
            const Frame &realNextFrame = *fit2;

            Frame predNextFrame;
            vector<const Frame *> history;
            history.push_back(&curFrame);
            model.AdvanceFrame(history, predNextFrame, params.lookAhead);

            FrameGrid realNextGrid(realNextFrame, params.gridScale);
            FrameGrid predNextGrid(predNextFrame, params.gridScale);

            for (int y = 0; y < realNextGrid.height() && y < predNextGrid.height(); y++) {
                for (int x = 0; x < realNextGrid.width() && x < predNextGrid.width(); x++) {
                    int realId = realNextGrid.getId(x, y);
                    int predId = predNextGrid.getId(x, y);

                    // Rejigger IDs in the case that either is empty
                    if (realId == -1) realId = numSprites - 1;
                    if (predId == -1) predId = numSprites - 1;

                    confusionMatrix(realId, predId) = confusionMatrix(realId, predId) + 1;
                    totalPredictions++;
                    if (realId == predId) totalCorrect++;
                }
            }

            ++fit1;
            ++fit2;
        }
        */
    }

    // Print out each confusion matrix
    for (int lookAhead = 1; lookAhead <= params.maxLookAhead; lookAhead++) {
        stringstream filename;
        filename << params.dirName << lookAhead << ".txt";
        ofstream os(filename.str());

        const Grid<int> &confusionMatrix = confusionMatrices[lookAhead - 1];

        // Print out the confusion matrix
        os << "Confusion matrix" << endl;

        // First print the column headers
        os << "\t";
        for (UINT i = 0; i < numSprites; i++) {
            os << spriteNames[i] << "\t";
        }
        os << endl;

        for (UINT i = 0; i < numSprites; i++) {
            os << spriteNames[i] << "\t";
            for (UINT j = 0; j < numSprites; j++) {
                os << confusionMatrix(i, j) << "\t";
            }
            os << endl;
        }
        int tc = totalCorrect[lookAhead - 1];
        int tp = totalPredictions[lookAhead - 1];
        double correctFraction = (tc + 0.0) / tp;
        os << "Number correct: " << tc << " / " << tp << " = " << correctFraction << endl;
        os << endl;
    }
}